{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c26ab996",
   "metadata": {},
   "source": [
    "# Self-RAG\n",
    "\n",
    "이 튜토리얼에서는 **Self-RAG**라는 Retrieval Augmented Generation(RAG) 전략을 소개하고, 이를 [LangGraph](https://langchain-ai.github.io/langgraph/)를 활용하여 구현하는 방법을 단계별로 설명합니다. \n",
    "\n",
    "Self-RAG는 검색된 문서와 생성된 응답에 대한 자기 반성(self-reflection) 및 자기 평가(self-evaluation)를 포함한 RAG 전략으로, RAG 기반 시스템의 성능 향상에 기여할 수 있습니다.\n",
    "\n",
    "![](./assets/langgraph-self-rag.png)\n",
    "\n",
    "---\n",
    "\n",
    "**Self-RAG란 무엇인가?**\n",
    "\n",
    "**Self-RAG**는 검색된 문서와 생성된 응답 모두에 대해 점검하고 검증하는 추가 단계를 포함하는 RAG 전략입니다.   \n",
    "전통적인 RAG에서는 검색된 정보를 기반으로 LLM이 답변을 생성하는 것이 주된 과정이었다면, Self-RAG에서는 **자체 평가**를 통해 다음과 같은 사항을 검증합니다.\n",
    "\n",
    "1. 검색할 필요성 판단: 현재 질문에 대해 추가 검색이 필요한지 여부를 판단합니다.\n",
    "2. 검색 결과 관련성 평가: 검색된 문서 조각(청크)이 질문 해결에 도움이 되는지 확인합니다.\n",
    "3. 응답 사실성 검증: 생성된 답변이 제공된 문서 청크에 의해 충분히 뒷받침되는지 평가합니다.\n",
    "4. 응답 품질 평가: 생성된 답변이 실제로 질문을 잘 해결하는지 측정합니다.\n",
    "\n",
    "이 과정을 통해 단순히 검색하고 답변을 생성하는 단계를 넘어, 생성된 응답의 품질과 사실성을 스스로 모니터링하고 향상시킬 수 있습니다.\n",
    "\n",
    "[Self-RAG 논문 바로가기](https://arxiv.org/abs/2310.11511)\n",
    "\n",
    "---\n",
    "\n",
    "**Self-RAG 주요 개념 정리**\n",
    "\n",
    "논문에서는 Self-RAG를 통해 다음과 같은 결정 과정을 제안합니다.\n",
    "\n",
    "1. **Retriever 사용 여부 결정**  \n",
    "   - 입력: `x (question)` 또는 `(x (question), y (generation))`  \n",
    "   - 출력: `yes, no, continue`  \n",
    "   이 단계는 추가 검색을 진행할지, 검색 없이 그대로 진행할지, 혹은 더 기다려볼지 결정합니다.\n",
    "   \n",
    "2. **관련성 평가 (Retrieval Grader)**  \n",
    "   - 입력: (`x (question)`, `d (chunk)`) for each `d` in `D`  \n",
    "   - 출력: `relevant` 또는 `irrelevant`  \n",
    "   검색된 문서 청크들이 실제로 질문에 답하는데 유용한 정보인지 판별합니다.\n",
    "   \n",
    "3. **사실성 검증 (Hallucination Grader)**  \n",
    "   - 입력: `x (question)`, `d (chunk)`, `y (generation)` for each `d` in `D`  \n",
    "   - 출력: `{fully supported, partially supported, no support}`  \n",
    "   생성된 응답이 검색 결과에 근거한 사실을 반영하는지, 혹은 환각(Hallucination)이 발생했는지 판단합니다.\n",
    "   \n",
    "4. **정답 품질 평가 (Answer Grader)**  \n",
    "   - 입력: `x (question)`, `y (generation)`  \n",
    "   - 출력: `{5, 4, 3, 2, 1}`  \n",
    "   생성된 응답이 질문을 어느 정도 해결하는지 점수화하여 평가합니다.\n",
    "\n",
    "---\n",
    "\n",
    "**이 튜토리얼에서 다룰 내용**\n",
    "\n",
    "이 튜토리얼은 LangGraph를 활용하여 Self-RAG 전략의 일부 아이디어를 구현하는 과정을 다룹니다.  \n",
    "다음과 같은 단계를 통해 Self-RAG 전략을 구축하고 실행하는 방법을 익히게 됩니다.\n",
    "\n",
    "- **Retriever**: 문서를 검색\n",
    "- **Retrieval Grader**: 검색된 문서의 관련성 평가\n",
    "- **Generate**: 질문에 대한 답변 생성\n",
    "- **Hallucination Grader**: 생성된 답변의 사실성(환각 여부) 검증\n",
    "- **Answer Grader**: 답변의 질문에 대한 관련성 평가\n",
    "- **Question Re-writer**: 쿼리 재작성\n",
    "- **그래프 생성 및 실행**: 정의한 노드로 그래프를 빌드하고 실행\n",
    "\n",
    "---\n",
    "\n",
    "**참고**\n",
    "\n",
    "- [LangGraph 공식 문서](https://langchain-ai.github.io/langgraph/)  \n",
    "- [Self-RAG 논문](https://arxiv.org/abs/2310.11511)  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de56f8df",
   "metadata": {},
   "source": [
    "## 환경설정"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6458235",
   "metadata": {},
   "outputs": [],
   "source": [
    "# API 키를 환경변수로 관리하기 위한 설정 파일\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "# API 키 정보 로드\n",
    "load_dotenv()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3043fdb5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# LangSmith 추적을 설정합니다. https://smith.langchain.com\n",
    "# !pip install -qU langchain-teddynote\n",
    "from langchain_teddynote import logging\n",
    "\n",
    "# 프로젝트 이름을 입력합니다.\n",
    "logging.langsmith(\"CH17-LangGraph-Use-Cases\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ec192a9",
   "metadata": {},
   "source": [
    "## 기본 PDF 기반 Retrieval Chain 생성\n",
    "\n",
    "여기서는 PDF 문서를 기반으로 Retrieval Chain 을 생성합니다. 가장 단순한 구조의 Retrieval Chain 입니다.\n",
    "\n",
    "단, LangGraph 에서는 Retriever 와 Chain 을 따로 생성합니다. 그래야 각 노드별로 세부 처리를 할 수 있습니다.\n",
    "\n",
    "**참고**\n",
    "\n",
    "- 이전 튜토리얼에서 다룬 내용이므로, 자세한 설명은 생략합니다.\n",
    "\n",
    "**실습에 활용한 문서**\n",
    "\n",
    "소프트웨어정책연구소(SPRi) - 2023년 12월호\n",
    "\n",
    "- 저자: 유재흥(AI정책연구실 책임연구원), 이지수(AI정책연구실 위촉연구원)\n",
    "- 링크: https://spri.kr/posts/view/23669\n",
    "- 파일명: `SPRI_AI_Brief_2023년12월호_F.pdf`\n",
    "\n",
    "_실습을 위해 다운로드 받은 파일을 `data` 폴더로 복사해 주시기 바랍니다_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "4cde195c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from rag.pdf import PDFRetrievalChain\n",
    "\n",
    "# PDF 문서를 로드합니다.\n",
    "pdf = PDFRetrievalChain([\"data/SPRI_AI_Brief_2023년12월호_F.pdf\"]).create_chain()\n",
    "\n",
    "# retriever와 chain을 생성합니다.\n",
    "pdf_retriever = pdf.retriever\n",
    "pdf_chain = pdf.chain"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65ec10e5",
   "metadata": {},
   "source": [
    "## 문서 검색 평가기 (Retrieval Grader)\n",
    "\n",
    "추후 retrieve 노드에서 문서에 대한 관련성 평가를 진행하기 위해 미리 정의합니다.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "5c85e4c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydantic import BaseModel, Field\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain_openai import ChatOpenAI\n",
    "from langchain_teddynote.models import get_model_name, LLMs\n",
    "\n",
    "# 최신모델 이름 설정\n",
    "MODEL_NAME = get_model_name(LLMs.GPT4o)\n",
    "\n",
    "\n",
    "# 데이터 모델 정의: 검색된 문서의 관련성을 이진 점수로 평가하기 위한 데이터 모델\n",
    "class GradeDocuments(BaseModel):\n",
    "    \"\"\"A binary score to determine the relevance of the retrieved documents.\"\"\"\n",
    "\n",
    "    # 문서가 질문에 관련이 있는지 여부를 'yes' 또는 'no'로 나타내는 필드\n",
    "    binary_score: str = Field(\n",
    "        description=\"Documents are relevant to the question, 'yes' or 'no'\"\n",
    "    )\n",
    "\n",
    "\n",
    "# LLM 초기화\n",
    "llm = ChatOpenAI(model=MODEL_NAME, temperature=0)\n",
    "\n",
    "# GradeDocuments 데이터 모델을 사용하여 LLM의 구조화된 출력 생성\n",
    "structured_llm_grader = llm.with_structured_output(GradeDocuments)\n",
    "\n",
    "# 시스템 프롬프트 정의: 검색된 문서가 사용자 질문에 관련이 있는지 평가하는 시스템 역할 정의\n",
    "system = \"\"\"You are a grader assessing relevance of a retrieved document to a user question. \\n \n",
    "    It does not need to be a stringent test. The goal is to filter out erroneous retrievals. \\n\n",
    "    If the document contains keyword(s) or semantic meaning related to the user question, grade it as relevant. \\n\n",
    "    Give a binary score 'yes' or 'no' score to indicate whether the document is relevant to the question.\"\"\"\n",
    "\n",
    "# 채팅 프롬프트 템플릿 생성\n",
    "grade_prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", system),\n",
    "        (\"human\", \"Retrieved document: \\n\\n {document} \\n\\n User question: {question}\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "# 검색 평가기 생성\n",
    "retrieval_grader = grade_prompt | structured_llm_grader"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c00a1d4",
   "metadata": {},
   "source": [
    "`retrieval_grader` 를 실행하여 검색된 문서의 관련성 평가를 수행합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2312e7a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 질문 정의\n",
    "question = \"삼성전자가 개발한 생성형 AI 의 이름은?\"\n",
    "\n",
    "# 문서 검색\n",
    "docs = pdf_retriever.invoke(question)\n",
    "\n",
    "# 검색된 문서 중 두 번째 문서의 페이지 콘텐츠 추출\n",
    "doc_txt = docs[1].page_content\n",
    "\n",
    "# 검색 평가기 호출 및 결과 출력\n",
    "print(retrieval_grader.invoke({\"question\": question, \"document\": doc_txt}))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "702b7f87",
   "metadata": {},
   "source": [
    "## 답변 생성 체인\n",
    "\n",
    "답변 생성 체인은 검색된 문서를 기반으로 답변을 생성하는 체인입니다.\n",
    "\n",
    "우리가 알고 있는 일반적인 Naive RAG 체인 입니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d080f695",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain import hub\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "\n",
    "# LangChain Hub에서 프롬프트 가져오기\n",
    "prompt = hub.pull(\"teddynote/rag-prompt\")\n",
    "\n",
    "# 기본 LLM 초기화, 모델 이름과 온도 설정\n",
    "llm = ChatOpenAI(model_name=MODEL_NAME, temperature=0)\n",
    "\n",
    "\n",
    "# 문서 포맷팅 함수\n",
    "def format_docs(docs):\n",
    "    return \"\\n\\n\".join(\n",
    "        [\n",
    "            f'<document><content>{doc.page_content}</content><source>{doc.metadata[\"source\"]}</source><page>{doc.metadata[\"page\"]+1}</page></document>'\n",
    "            for doc in docs\n",
    "        ]\n",
    "    )\n",
    "\n",
    "\n",
    "# RAG 체인 생성\n",
    "rag_chain = prompt | llm | StrOutputParser()\n",
    "\n",
    "# 체인 실행\n",
    "generation = rag_chain.invoke({\"context\": format_docs(docs), \"question\": question})\n",
    "print(generation)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41d4050d",
   "metadata": {},
   "source": [
    "## 답변의 할루시네이션 여부를 평가\n",
    "\n",
    "`groundedness_grader` 를 생성하고 생성된 답변과 `context` 를 기반하여 답변의 할루시네이션 평가를 진행합니다.\n",
    "\n",
    "\n",
    "`yes` 인 경우 답변의 할루시네이션이 없을을 의미합니다. 반대로, `no` 인 경우 답변이 할루시네이션으라고 간주합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "2f46b9ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydantic import BaseModel, Field\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "\n",
    "# 데이터 모델 정의: 생성된 답변이 사실에 기반하고 있는지 여부를 이진 점수로 평가하기 위한 데이터 모델\n",
    "class Groundednesss(BaseModel):\n",
    "    \"\"\"A binary score indicating whether the generated answer is grounded in the facts.\"\"\"\n",
    "\n",
    "    # 답변이 사실에 기반하고 있는지 여부를 'yes' 또는 'no'로 나타내는 필드\n",
    "    binary_score: str = Field(\n",
    "        description=\"Answer is grounded in the facts, 'yes' or 'no'\"\n",
    "    )\n",
    "\n",
    "\n",
    "# LLM 초기화\n",
    "llm = ChatOpenAI(model=MODEL_NAME, temperature=0)\n",
    "\n",
    "# 구조화된 출력과 함께 LLM 설정\n",
    "structured_llm_grader = llm.with_structured_output(Groundednesss)\n",
    "\n",
    "# 시스템 프롬프트 정의\n",
    "system = \"\"\"You are a grader assessing whether an LLM generation is grounded in / supported by a set of retrieved facts. \\n \n",
    "Give a binary score 'yes' or 'no'. 'Yes' means that the answer is grounded in / supported by the set of facts.\"\"\"\n",
    "\n",
    "# 채팅 프롬프트 템플릿 생성\n",
    "groundedness_checking_prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", system),\n",
    "        (\"human\", \"Set of facts: \\n\\n {documents} \\n\\n LLM generation: {generation}\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "# 답변의 할루시네이션 평가기 생성\n",
    "groundedness_grader = groundedness_checking_prompt | structured_llm_grader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc4f276a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 할루시네이션 평가기 호출(yes: 사실에 기반하고 있음, no: 사실에 기반하고 있지 않음)\n",
    "groundedness_grader.invoke({\"documents\": format_docs(docs), \"generation\": generation})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54c773d2",
   "metadata": {},
   "source": [
    "## 답변의 관련성 평가\n",
    "\n",
    "생성된 답변이 질문(question) 과의 관련성 있는 답변인지 여부를 평가합니다.\n",
    "\n",
    "`yes` 인 경우는 관련성 있음. `no` 인 경우는 관련성 없음을 의미합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "37e10188",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydantic import BaseModel, Field\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "\n",
    "class GradeAnswer(BaseModel):\n",
    "    \"\"\"A binary score indicating whether the question is addressed.\"\"\"\n",
    "\n",
    "    # 답변의 관련성 평가: 'yes' 또는 'no'로 표기(yes: 관련성 있음, no: 관련성 없음)\n",
    "    binary_score: str = Field(\n",
    "        description=\"Answer addresses the question, 'yes' or 'no'\"\n",
    "    )\n",
    "\n",
    "\n",
    "llm = ChatOpenAI(model=MODEL_NAME, temperature=0)\n",
    "\n",
    "# llm 에 GradeAnswer 바인딩\n",
    "structured_llm_grader = llm.with_structured_output(GradeAnswer)\n",
    "\n",
    "# 시스템 프롬프트 정의\n",
    "system = \"\"\"You are a grader assessing whether an answer addresses / resolves a question \\n \n",
    "     Give a binary score 'yes' or 'no'. Yes' means that the answer resolves the question.\"\"\"\n",
    "\n",
    "# 프롬프트 생성\n",
    "answer_grader_prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", system),\n",
    "        (\"human\", \"User question: \\n\\n {question} \\n\\n LLM generation: {generation}\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "# 답변 평가기 생성\n",
    "answer_grader = answer_grader_prompt | structured_llm_grader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07c593ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 답변 평가기 호출(yes: 질문을 해결함, no: 질문을 해결하지 않음)\n",
    "answer_grader.invoke({\"question\": question, \"generation\": generation})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd05275c",
   "metadata": {},
   "source": [
    "## 질문 재작성기(Question Rewriter)\n",
    "\n",
    "사용자가 입력한 질문 재작성기를 생성합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "806f6d04",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain_openai import ChatOpenAI\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "\n",
    "\n",
    "llm = ChatOpenAI(model=MODEL_NAME, temperature=0)\n",
    "\n",
    "# 시스템 프롬프트 정의\n",
    "# 입력 질문을 벡터스토어 검색에 최적화된 형태로 변환하는 시스템 역할 정의\n",
    "system = \"\"\"You a question re-writer that converts an input question to a better version that is optimized \\n \n",
    "     for vectorstore retrieval. Look at the input and try to reason about the underlying semantic intent / meaning.\"\"\"\n",
    "\n",
    "# 시스템 메시지와 초기 질문을 포함한 프롬프트 템플릿 생성\n",
    "re_write_prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", system),\n",
    "        (\n",
    "            \"human\",\n",
    "            \"Here is the initial question: \\n\\n {question} \\n Formulate an improved question.\",\n",
    "        ),\n",
    "    ]\n",
    ")\n",
    "\n",
    "# 질문 재작성기 생성\n",
    "question_rewriter = re_write_prompt | llm | StrOutputParser()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e681966",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 질문 재작성기 호출\n",
    "question_rewriter.invoke({\"question\": question})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f08d720d",
   "metadata": {},
   "source": [
    "## 상태 정의\n",
    "\n",
    "상태(State)를 정의합니다.\n",
    "\n",
    "- `question`: 사용자가 입력한 질문\n",
    "- `generation`: 생성된 응답\n",
    "- `documents`: 검색된 문서 목록"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "52775711",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "from typing_extensions import TypedDict, Annotated\n",
    "\n",
    "\n",
    "# 그래프의 상태를 나타내는 클래스 정의\n",
    "class GraphState(TypedDict):\n",
    "    # 질문을 나타내는 문자열\n",
    "    question: Annotated[str, \"Question\"]\n",
    "    # LLM에 의해 생성된 응답을 나타내는 문자열\n",
    "    generation: Annotated[str, \"LLM Generation\"]\n",
    "    # 문서의 목록을 나타내는 문자열 리스트\n",
    "    documents: Annotated[List[str], \"Retrieved Documents\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "78a0836f",
   "metadata": {},
   "source": [
    "## 노드 정의\n",
    "\n",
    "- `retrieve`: 문서 검색\n",
    "- `grade_documents`: 문서 평가\n",
    "- `generate`: 답변 생성\n",
    "- `transform_query`: 질문 재작성\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "4cbd12b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 문서 검색\n",
    "def retrieve(state):\n",
    "    print(\"==== [RETRIEVE] ====\")\n",
    "    question = state[\"question\"]\n",
    "\n",
    "    # 검색 수행\n",
    "    documents = pdf_retriever.invoke(question)\n",
    "    return {\"documents\": documents}\n",
    "\n",
    "\n",
    "# 답변 생성\n",
    "def generate(state):\n",
    "    print(\"==== [GENERATE] ====\")\n",
    "    question = state[\"question\"]\n",
    "    documents = state[\"documents\"]\n",
    "\n",
    "    # RAG 생성\n",
    "    generation = rag_chain.invoke({\"context\": documents, \"question\": question})\n",
    "    return {\"generation\": generation}\n",
    "\n",
    "\n",
    "# 검색된 문서의 관련성 평가\n",
    "def grade_documents(state):\n",
    "    print(\"==== [GRADE DOCUMENTS] ====\")\n",
    "    question = state[\"question\"]\n",
    "    documents = state[\"documents\"]\n",
    "\n",
    "    # 각 문서 점수 평가\n",
    "    filtered_docs = []\n",
    "    for d in documents:\n",
    "        score = retrieval_grader.invoke(\n",
    "            {\"question\": question, \"document\": d.page_content}\n",
    "        )\n",
    "        grade = score.binary_score\n",
    "        if grade == \"yes\":\n",
    "            print(\"==== GRADE: DOCUMENT RELEVANT ====\")\n",
    "            filtered_docs.append(d)\n",
    "        else:\n",
    "            print(\"==== GRADE: DOCUMENT NOT RELEVANT ====\")\n",
    "            continue\n",
    "    return {\"documents\": filtered_docs}\n",
    "\n",
    "\n",
    "# 질문 변환\n",
    "def transform_query(state):\n",
    "    print(\"==== [TRANSFORM QUERY] ====\")\n",
    "    question = state[\"question\"]\n",
    "\n",
    "    # 질문 재작성\n",
    "    better_question = question_rewriter.invoke({\"question\": question})\n",
    "    return {\"question\": better_question}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ef3da3f",
   "metadata": {},
   "source": [
    "## 조건부 엣지 정의\n",
    "\n",
    "`decide_to_generate` 함수는 검색된 문서의 관련성 평가 결과에 따라 답변 생성 여부를 결정합니다.\n",
    "\n",
    "`grade_generation_v_documents_and_question` 함수는 생성된 답변의 문서 및 질문과의 관련성 평가 결과에 따라 답변 생성 여부를 결정합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "87d1f724",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 답변 생성 여부 결정\n",
    "def decide_to_generate(state):\n",
    "    print(\"==== [ASSESS GRADED DOCUMENTS] ====\")\n",
    "    state[\"question\"]\n",
    "    filtered_documents = state[\"documents\"]\n",
    "\n",
    "    if not filtered_documents:\n",
    "        # 모든 문서가 관련성이 없는 경우\n",
    "        # 새로운 쿼리 생성\n",
    "        print(\n",
    "            \"==== [DECISION: ALL DOCUMENTS ARE NOT RELEVANT TO QUESTION, TRANSFORM QUERY] ====\"\n",
    "        )\n",
    "        return \"transform_query\"\n",
    "    else:\n",
    "        # 관련 문서가 있는 경우 답변 생성\n",
    "        print(\"==== [DECISION: GENERATE] ====\")\n",
    "        return \"generate\"\n",
    "\n",
    "\n",
    "# 생성된 답변의 문서 및 질문과의 관련성 평가\n",
    "def grade_generation_v_documents_and_question(state):\n",
    "    print(\"==== [CHECK HALLUCINATIONS] ====\")\n",
    "    question = state[\"question\"]\n",
    "    documents = state[\"documents\"]\n",
    "    generation = state[\"generation\"]\n",
    "\n",
    "    score = groundedness_grader.invoke(\n",
    "        {\"documents\": documents, \"generation\": generation}\n",
    "    )\n",
    "    grade = score.binary_score\n",
    "\n",
    "    # 환각 여부 확인\n",
    "    if grade == \"yes\":\n",
    "        print(\"==== [DECISION: GENERATION IS GROUNDED IN DOCUMENTS] ====\")\n",
    "        # 질문 해결 여부 확인\n",
    "        print(\"==== [GRADE GENERATION vs QUESTION] ====\")\n",
    "        score = answer_grader.invoke({\"question\": question, \"generation\": generation})\n",
    "        grade = score.binary_score\n",
    "        if grade == \"yes\":\n",
    "            print(\"==== [DECISION: GENERATION ADDRESSES QUESTION] ====\")\n",
    "            return \"relevant\"\n",
    "        else:\n",
    "            print(\"==== [DECISION: GENERATION DOES NOT ADDRESS QUESTION] ====\")\n",
    "            return \"not relevant\"\n",
    "    else:\n",
    "        print(\"==== [DECISION: GENERATION IS NOT GROUNDED IN DOCUMENTS, RE-TRY] ====\")\n",
    "        return \"hallucination\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee9a8d0e",
   "metadata": {},
   "source": [
    "## 그래프 생성\n",
    "\n",
    "이전에 작성한 노드와 엣지를 통해 그래프를 생성합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "204852a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.graph import END, StateGraph, START\n",
    "from langgraph.checkpoint.memory import MemorySaver\n",
    "\n",
    "# 그래프 상태 초기화\n",
    "workflow = StateGraph(GraphState)\n",
    "\n",
    "# 노드 정의\n",
    "workflow.add_node(\"retrieve\", retrieve)  # retrieve\n",
    "workflow.add_node(\"grade_documents\", grade_documents)  # grade documents\n",
    "workflow.add_node(\"generate\", generate)  # generatae\n",
    "workflow.add_node(\"transform_query\", transform_query)  # transform_query\n",
    "\n",
    "# 엣지 정의\n",
    "workflow.add_edge(START, \"retrieve\")\n",
    "workflow.add_edge(\"retrieve\", \"grade_documents\")\n",
    "\n",
    "# 문서 평가 노드에서 조건부 엣지 추가\n",
    "workflow.add_conditional_edges(\n",
    "    \"grade_documents\",\n",
    "    decide_to_generate,\n",
    "    {\n",
    "        \"transform_query\": \"transform_query\",\n",
    "        \"generate\": \"generate\",\n",
    "    },\n",
    ")\n",
    "\n",
    "# 엣지 정의\n",
    "workflow.add_edge(\"transform_query\", \"retrieve\")\n",
    "\n",
    "# 답변 생성 노드에서 조건부 엣지 추가\n",
    "workflow.add_conditional_edges(\n",
    "    \"generate\",\n",
    "    grade_generation_v_documents_and_question,\n",
    "    {\n",
    "        \"hallucination\": \"generate\",\n",
    "        \"relevant\": END,\n",
    "        \"not relevant\": \"transform_query\",\n",
    "    },\n",
    ")\n",
    "\n",
    "# 그래프 컴파일\n",
    "app = workflow.compile(checkpointer=MemorySaver())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1fd1bad",
   "metadata": {},
   "source": [
    "그래프를 시각화 합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc3b65ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_teddynote.graphs import visualize_graph\n",
    "\n",
    "visualize_graph(app)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d8c305a",
   "metadata": {},
   "source": [
    "## 그래프 실행\n",
    "\n",
    "생성한 그래프를 실행합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a6458fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.runnables import RunnableConfig\n",
    "from langchain_teddynote.messages import stream_graph, invoke_graph, random_uuid\n",
    "\n",
    "# config 설정(재귀 최대 횟수, thread_id)\n",
    "config = RunnableConfig(recursion_limit=10, configurable={\"thread_id\": random_uuid()})\n",
    "\n",
    "# 질문 입력\n",
    "inputs = {\n",
    "    \"question\": \"삼성전자가 개발한 생성형 AI 의 이름은?\",\n",
    "}\n",
    "\n",
    "# 그래프 실행\n",
    "invoke_graph(\n",
    "    app, inputs, config, [\"retrieve\", \"transform_query\", \"grade_documents\", \"generate\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "218f12e3",
   "metadata": {},
   "source": [
    "만약, 사용자의 질문에 대한 관련성 평가에서 지속적인 failure가 발생하는 경우, 다음과 같이 재귀상태로 빠질 수 있습니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20b4db9c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.errors import GraphRecursionError\n",
    "\n",
    "# config 설정(재귀 최대 횟수, thread_id)\n",
    "config = RunnableConfig(recursion_limit=10, configurable={\"thread_id\": random_uuid()})\n",
    "\n",
    "# 질문 입력\n",
    "inputs = {\n",
    "    \"question\": \"테디노트가 개발한 생성형 AI 의 이름은?\",\n",
    "}\n",
    "\n",
    "try:\n",
    "    # 그래프 실행\n",
    "    stream_graph(\n",
    "        app,\n",
    "        inputs,\n",
    "        config,\n",
    "        [\"retrieve\", \"transform_query\", \"grade_documents\", \"generate\"],\n",
    "    )\n",
    "except GraphRecursionError as recursion_error:\n",
    "    print(f\"GraphRecursionError: {recursion_error}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2cda1a04",
   "metadata": {},
   "source": [
    "이렇게 재귀상태에 빠지지 않도록 탈출을 위한 로직 수정(흐름의 수정)이 필요합니다."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "langchain-kr-lwwSZlnu-py3.11",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
